
import numpy as np

class DHAAggregator:
    """动态分层注意力聚合框架"""
    def __init__(self, entropy_threshold=0.7):
        self.entropy_threshold = entropy_threshold
        self.temp_controller = TemperatureAttentionController()
        
    def device_state_entropy(self, hetero_score, param_entropy, beta=0.6):
        """设备状态熵融合定义 (式9)"""
        return beta * hetero_score + (1 - beta) * param_entropy
    
    def dynamic_hierarchization(self, state_entropy, temperature):
        """熵阈驱动的动态分层 (式10)"""
        # 动态调整阈值
        dynamic_threshold = self.entropy_threshold * (1 + 0.2 * (1 - temperature))
        
        if state_entropy < dynamic_threshold * 0.8:
            return "high"  # 高资源层
        elif state_entropy < dynamic_threshold:
            return "medium"  # 中资源层
        else:
            return "low"  # 低资源层
    
    def attention_aggregation(self, layer_type, entropies, temperature):
        """层内熵驱动加权 (3.3节)"""
        if layer_type == "high":
            weights = np.exp(-entropies / temperature)
        else:  # 中低资源层
            weights = np.ones_like(entropies)
            weights[entropies > self.entropy_threshold] = 0  # 过滤高熵设备
        return weights / np.sum(weights)
    
    def lightweight_correction(self, teacher_grad, student_grad, temperature):
        """轻量化梯度校正 (式11)"""
        kl_div = np.sum(teacher_grad * np.log(teacher_grad/student_grad + 1e-10))
        correction_strength = 1 - temperature  # 温度越低约束越强
        return student_grad + correction_strength * kl_div * (teacher_grad - student_grad)
